import numpy as np
from sklearn.manifold import TSNE
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import time
import json
import pdb

meta_type = 1
shot = 10

t0 = time.time()
file1 = 'model_path/pooled_feat/score_gt_%.0f_%.0f_test_gt.json' % (meta_type, shot)
data1 = json.load(open(file1, 'r'))
cls_vector_all = np.array(data1['cls_vector'])
cls_gt_all = np.array(data1['cls_gt'])
pooled_feat_all = np.array(data1['pooled_feat'])
feature_num = pooled_feat_all.shape[0]

'''
file2 = 'model_path/pooled_feat/score_gt_%.0f_%.0f_margin_trainval_gt.json' % (meta_type, shot)
data2 = json.load(open(file2, 'r'))
cls_vector_all = np.vstack((cls_vector_all, np.array(data2['cls_vector'])))
cls_gt_all = np.vstack((cls_gt_all, np.array(data2['cls_gt'])))
pooled_feat_all = np.vstack((pooled_feat_all, np.array(data2['pooled_feat'])))
'''


t1 = time.time()
print('After %.2f seconds, loaded the whole data.' % (t1 - t0))
print('tsne start')
tsne_all = TSNE(n_components = 2).fit_transform(pooled_feat_all)
t2 = time.time()
print('tsne time used: %.2f' % float(t2 - t1))

tsne_data = {'tsne_data': tsne_all.tolist(),
            'cls_vector': cls_vector_all.tolist(),
            'cls_gt': cls_gt_all.tolist()}
file3 = 'model_path/pooled_feat/tsne_data_test_gt.json'
with open(file3,'w') as f:
    json.dump(tsne_data, f)

xmin = tsne_all[:, 0].min()
xmax = tsne_all[:, 0].max()
ymin = tsne_all[:, 1].min()
ymax = tsne_all[:, 1].max()

color = ['#66FFFF',
'#808080',
'#00FFFF',
'#7FFFD4',
'#008000',
'#90EE90',
'#DC143C',
'#FFA500',
'#FFEBCD',
'#0000FF',
'#8A2BE2',
'#A52A2A',
'#DEB887',
'#5F9EA0',
'#FF7F50',
'#D2691E',
'#7FFF00',
'#6495ED',
'#000000',
'#D3D3D3',
'#FFD700']

tsne = tsne_all[ : feature_num, :]
cls_gt = cls_gt_all[:feature_num, :]
cls_vector = cls_vector_all[:feature_num, :]

for score in range(10, 0, -1):
    for cls in range(1, 21):
    
        keep = cls_gt[:, 0] == cls
        x = tsne[keep, 0]
        y = tsne[keep, 1]
        cls_score = cls_vector[keep, :]

        keep = (cls_score[:, cls] >= (score-1)/10) & (cls_score[:, cls] < score/10)
        
        x_score = x[keep]
        y_score = y[keep]
        if cls < 16:
            plt.scatter(x_score, y_score, s=1, c = color[cls], alpha = 1)
        else:
            plt.scatter(x_score, y_score, s=1, marker='x', c = color[cls], alpha = 1)

t3 = time.time()
print('scatter plot time used: %.2f' % (t3 - t2))
if meta_type == 1:
    plt.legend(['aeroplane', 'bicycle', 'boat', 'bottle', 'car', 'cat', 'chair', 'diningtable', 'dog', 'horse', 'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'bird', 'bus', 'cow', 'motorbike', 'sofa'], bbox_to_anchor=(1, 1))
elif meta_type == 2:
    plt.legend(['bicycle', 'bird', 'boat', 'bus', 'car', 'cat', 'chair', 'diningtable', 'dog', 'motorbike',
                         'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'aeroplane', 'bottle', 'cow',
                         'horse', 'sofa'], bbox_to_anchor=(1, 1))
elif meta_type == 3:
    plt.legend(['aeroplane', 'bicycle', 'bird', 'bottle', 'bus', 'car', 'chair', 'cow', 'diningtable',
                         'dog', 'horse', 'person', 'pottedplant', 'train', 'tvmonitor', 'boat', 'cat', 'motorbike',
                         'sheep', 'sofa'], bbox_to_anchor=(1, 1))
      
plt.xlim(xmin, xmax)
plt.ylim(ymin, ymax)
plt.tight_layout()
plt.savefig('tsne_test_gt.jpg', dpi=1000)
plt.cla()
'''
for score in range(1, 6):
    for cls in range(1, 21):
    
        keep = cls_gt[:, 0] == cls
        x = tsne[keep, 0]
        y = tsne[keep, 1]
        cls_score = cls_vector[keep, :]

        keep = (cls_score[:, cls] >= (score-1)/10) & (cls_score[:, cls] < score/10)
        
        x_score = x[keep]
        y_score = y[keep]
        if cls < 16:
            plt.scatter(x_score, y_score, s=1, c = color[cls], alpha = 1)
        else:
            plt.scatter(x_score, y_score, s=1, marker='x', c = color[cls], alpha = 1)
    if score == 5:
        plt.savefig('tsne_score_0_%.1f_cag.jpg' % (score/10), dpi=1000)
'''
'''
tsne = tsne_all[feature_num:, :]
cls_gt = cls_gt_all[feature_num:, :]
cls_vector = cls_vector_all[feature_num:, :]

for score in range(10, 0, -1):
    for cls in range(1, 21):
    
        keep = cls_gt[:, 0] == cls
        x = tsne[keep, 0]
        y = tsne[keep, 1]
        cls_score = cls_vector[keep, :]

        keep = (cls_score[:, cls] >= (score-1)/10) & (cls_score[:, cls] < score/10)
        
        x_score = x[keep]
        y_score = y[keep]
        if cls < 16:
            plt.scatter(x_score, y_score, s=1, c = color[cls], alpha = 1)
        else:
            plt.scatter(x_score, y_score, s=1, marker='x', c = color[cls], alpha = 1)
    if score == 5:
        plt.savefig('tsne_score_%.1f_10_margin.jpg' % (score/10), dpi=1000)

t3 = time.time()
print('scatter plot time used: %.2f' % (t3 - t2))
if meta_type == 1:
    plt.legend(['aeroplane', 'bicycle', 'boat', 'bottle', 'car', 'cat', 'chair', 'diningtable', 'dog', 'horse', 'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'bird', 'bus', 'cow', 'motorbike', 'sofa'], bbox_to_anchor=(0.95, 1))
elif meta_type == 2:
    plt.legend(['bicycle', 'bird', 'boat', 'bus', 'car', 'cat', 'chair', 'diningtable', 'dog', 'motorbike',
                         'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'aeroplane', 'bottle', 'cow',
                         'horse', 'sofa'], bbox_to_anchor=(0.95, 1))
elif meta_type == 3:
    plt.legend(['aeroplane', 'bicycle', 'bird', 'bottle', 'bus', 'car', 'chair', 'cow', 'diningtable',
                         'dog', 'horse', 'person', 'pottedplant', 'train', 'tvmonitor', 'boat', 'cat', 'motorbike',
                         'sheep', 'sofa'], bbox_to_anchor=(0.95, 1))
      
plt.xlim(xmin, xmax)
plt.ylim(ymin, ymax)
plt.savefig('tsne_margin.jpg', dpi=1000)
plt.cla()
for score in range(1, 6):
    for cls in range(1, 21):
    
        keep = cls_gt[:, 0] == cls
        x = tsne[keep, 0]
        y = tsne[keep, 1]
        cls_score = cls_vector[keep, :]

        keep = (cls_score[:, cls] >= (score-1)/10) & (cls_score[:, cls] < score/10)
        
        x_score = x[keep]
        y_score = y[keep]
        if cls < 16:
            plt.scatter(x_score, y_score, s=1, c = color[cls], alpha = 1)
        else:
            plt.scatter(x_score, y_score, s=1, marker='x', c = color[cls], alpha = 1)
    if score == 5:
        plt.savefig('tsne_score_0_%.1f_margin.jpg' % (score/10), dpi=1000)
'''